package com.liuhengjia.aiknowledgebase.service.impl;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.lark.oapi.service.drive.v1.model.File;
import com.liuhengjia.aiknowledgebase.entity.BaseDoc;
import com.liuhengjia.aiknowledgebase.service.BaseDocService;
import com.liuhengjia.aiknowledgebase.api.response.DocVO;
import com.liuhengjia.aiknowledgebase.third.FeiShuService;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import lombok.RequiredArgsConstructor;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

/**
 * 文档操作服务层实现类
 */
@Service
@RequiredArgsConstructor
public class DocServiceImpl implements com.liuhengjia.aiknowledgebase.service.DocService {
    private final FeiShuService feiShuService;
    private final BaseDocService baseDocService;
    private final VectorStore vectorStore;

    @Override
    @Transactional(rollbackFor = Exception.class)
    public void loadData() {
        List<File> fileList = feiShuService.listDocs();
        List<BaseDoc> baseDocList = baseDocService.list();

        Map<String, BaseDoc> stringBaseDocMap = baseDocList.stream().collect(Collectors.toMap(BaseDoc::getDocId, doc -> doc));

        List<BaseDoc> newDocs = fileList.stream()
                .filter(v -> !stringBaseDocMap.containsKey(v.getToken()))
                .map(v -> BaseDoc.builder().docId(v.getToken()).docName(v.getName()).url(v.getUrl()).modifiedTime(v.getModifiedTime()).build())
                .toList();
        addToVectorStore(newDocs);

        List<String> fileIdList = fileList.stream()
                .map(File::getToken)
                .toList();
        List<BaseDoc> deletedDocList = baseDocList.stream()
                .filter(v -> !fileIdList.contains(v.getDocId()))
                .toList();
        if (!CollectionUtils.isEmpty(deletedDocList)) {
            removeFromVectorStore(deletedDocList);
        }

        List<BaseDoc> updatedDocs = fileList.stream()
                .filter(v -> stringBaseDocMap.containsKey(v.getToken()))
                .filter(v -> !v.getModifiedTime().equals(stringBaseDocMap.get(v.getToken()).getModifiedTime()))
                .map(v -> {
                    BaseDoc baseDoc = stringBaseDocMap.get(v.getToken());
                    return BaseDoc.builder().id(baseDoc.getId()).docId(baseDoc.getDocId()).docName(v.getName()).url(v.getUrl()).vectorDocId(baseDoc.getVectorDocId()).modifiedTime(v.getModifiedTime()).build();
                })
                .toList();
        updateVectorStore(updatedDocs);
    }

    @Override
    public List<DocVO> list() {
        return baseDocService.list().stream()
                .map(v -> DocVO.builder().docId(v.getDocId()).name(v.getDocName()).url(v.getUrl()).build())
                .toList();
    }

    /**
     * 添加至向量数据库
     *
     * @param baseDocList 文档实体对象列表
     */
    private void addToVectorStore(List<BaseDoc> baseDocList) {
        baseDocList.forEach(e -> {
            Document document = new Document(feiShuService.readDoc(e.getDocId()), new HashMap<>());
            vectorStore.write(List.of(document));
            e.setVectorDocId(document.getId());
        });
        baseDocService.saveBatch(baseDocList);
    }

    /**
     * 更新至向量数据库
     *
     * @param baseDocList 文档实体对象列表
     */
    private void updateVectorStore(List<BaseDoc> baseDocList) {
        removeFromVectorStore(baseDocList);
        addToVectorStore(baseDocList);
        baseDocService.updateBatchById(baseDocList);
    }

    /**
     * 删除向量数据库
     *
     * @param baseDocList 文档实体对象列表
     */
    private void removeFromVectorStore(List<BaseDoc> baseDocList) {
        if (CollectionUtils.isEmpty(baseDocList)) {
            return;
        }
        vectorStore.delete(baseDocList.stream().map(BaseDoc::getVectorDocId).toList());
        baseDocService.remove(new LambdaQueryWrapper<BaseDoc>().in(BaseDoc::getDocId, baseDocList.stream().map(BaseDoc::getDocId).toList()));
    }
}
